import gevent
import gevent.socket
import gevent.server

import uuid

from collections import defaultdict
from cStringIO import StringIO
from mimetools import Message

__version__ = '0.1.0'

BUFF_LEN = 8192
CONNECT = 'CONNECT'
OTHER = set(('OPTIONS', 'GET', 'HEAD', 'POST', 'PUT',
                             'DELETE', 'TRACE'))
HTTPVER = 'HTTP/1.1'
VERSION = 'juru Proxy/'+__version__

session_connections = defaultdict(dict)

class ConnectionClosed(gevent.socket.error):
    pass

def send(sock, data):
    length = len(data)
    sent = sock.send(data)
    while sent < length:
        sent += sock.send(data[sent:])

def get_connection_ex(sid, hostport, host):
    global session_connections
    connections = session_connections[sid]
#     if host in connections:
#         conn = connections[host]
#         try:
#             conn.getpeername() # is closed?
#             return conn
#         except gevent.socket.error:
#             pass
    sock = gevent.socket.socket()
    sock.connect(hostport)
    print 'connected:', hostport
#    connections[host] = sock
    return sock

def get_connection(sid, hostport):
    global session_connections
    connections = session_connections[sid]
#     if hostport in connections:
#         conn = connections[hostport]
#         try:
#             conn.getpeername() # is closed?
#             return conn
#         except gevent.socket.error:
#             pass
    if ':' in hostport:
        host, port = hostport.split(':')
        port = int(port)
    else:
        host = hostport
        port = 80
    sock = gevent.socket.socket()
    print 'connect:', hostport
    sock.connect((host, port))
    print 'connected:', hostport
#    connections[hostport] = sock
    return sock

def get_one_from_buff(buff):
    i = buff.find('\r\n\r\n')
    # <CR><LF> not in buff
    if i < 0:
        return None
    sio = StringIO(buff)
    first_line = sio.readline()
    headers = Message(sio)
    t = sio.read()
    print first_line
    print 'get_one_from_buff, buff_len:', len(t)
    return first_line, headers, t

def get_one(sock, buff):
    while True:
        while  True:
            request = get_one_from_buff(buff)
            if request:
                print 'get_one, buff_len:', len(buff)
                return request
            else:
                break
        t = sock.recv(BUFF_LEN)
        if not t:
            raise ConnectionClosed()
        buff += t

def forword(from_, to, bytes):
    print 'call forword.'
    while 1:
        buff = from_.recv(BUFF_LEN)
        if not buff:
            raise gevent.socket.error()
        buff_len = len(buff)
        print bytes, 'sent:', buff_len
        if buff_len >= bytes:
            send(to, buff[:bytes])
            print 'forword returned.'
            return buff[bytes:]
        send(to, buff)
        bytes -= buff_len

def _read_chunk_length(sock, buff):
    while '\r\n' not in buff:
        t = sock.recv(BUFF_LEN)
        if not t:
            raise gevent.socket.error('_read_chunk_length error.')
        buff += t
    length, buff = buff.split('\r\n', 1)
    return int(length.strip(), 16), buff

def _send_footer(from_, to, buff):
    while True:
        if buff.startswith('\r\n'):
            send(to, '\r\n')
            return buff[2:]
        if '\r\n\r\n' in buff:
            t, buff = buff.split('\r\n\r\n', 1)
            send(to, t + '\r\n\r\n')
            return buff
        t = from_.recv(BUFF_LEN)
        if not t:
            raise gevent.socket.error('send_footer error.')
        buff += t

def _chunked_send(from_, to, buff):
    while True:
        length, buff = _read_chunk_length(from_, buff)
        if length == 0:
            send(to, '0\r\n') # send 'end'
            return _send_footer(from_, to, buff)
        send(to, hex(length) + '\r\n')
        buff_len = len(buff)
        _len = length + 2 # +2 for '\r\n'
        if buff_len >= _len:
            send(to, buff[:_len])
            buff = buff[_len:]
            continue
        else:
            print buff_len, _len, _len - buff_len
            send(to, buff)
            buff = forword(from_, to, _len - buff_len)

def send_response_to_client(server, client):
    buff = ''
    ps, headers, buff = get_one(server, buff) # ps = protocol and status
    headers_raw_data = ''.join(headers.headers)
    cl = headers.get('Content-Length')
    print 'back:', repr(ps), cl
    print headers_raw_data
    if not cl:
        send(client, ps + headers_raw_data + '\r\n')
        te = headers.get('Transfer-Encoding')
        assert te
        if 'chunked' not in te.lower():
            raise gevent.socket.error('no Content-Length, no Transfer-Encoding.')
        _chunked_send(server, client, buff)
        return 
    cl = int(cl)
    buff_len = len(buff)
    print cl, buff_len
    if buff_len >= cl:
        send(client, ps + headers_raw_data + '\r\n' + buff[:cl])
        return 
    send(client, ps + headers_raw_data + '\r\n' + buff)
    forword(server, client, cl - buff_len)

def _process_connect(sid, mpp, headers, client):
    print '$' * 30
    print mpp
    print ''.join(headers.headers)
    method, path, protocol = mpp.split()
    host, port = path.split(':')
    port = int(port)
    server = get_connection_ex(sid, (host, port), headers.get('Host'))
    send(client, ('HTTP/1.1 200 Connection established\r\n'
        'Proxy-agent:juru\r\n\r\n'))
    import gevent.ssl
    client = gevent.ssl.wrap_socket(client)
    print '$' * 30

def proxy_request(sid, mpp, headers, buff, client):
    if mpp.startswith(CONNECT):
        _process_connect(sid, mpp, headers, client)
        return buff
    host = headers.get('Host')
    print 'host:', host
    assert host
    server = get_connection(sid, host)
    print '8'*30
    print server.getpeername()
    cl = headers.get('Content-Length')
    print 'connect-length:', cl
    headers_raw_data = ''.join(headers.headers)
#   print headers_raw_data
    if not cl:
        send(server, mpp + headers_raw_data + '\r\n')
        send_response_to_client(server, client)
        return buff
    cl = int(cl)
    buff_len = len(buff)
    if buff_len >= cl:
        send(server, mpp + headers_raw_data + '\r\n' + buff[:cl])
        buff = buff[cl:]
        send_response_to_client(server, client)
        return buff
    send(server, mpp + headers_raw_data + '\r\n' + buff)
    buff = forword(client, server, cl - buff_len)
    send_response_to_client(server, client)
    return buff

def handler(sock, addr):
    global session_connections
    session_id = uuid.uuid4()
    buff = ''
    while  True:
        print addr, 'start', '*' * 30
        try:
            request = get_one(sock, buff)
        except ConnectionClosed:
            request = None
        if not request:
            break
        mpp, headers, buff = request
        print addr, mpp
        buff = proxy_request(session_id, mpp, headers, buff, sock)
        print addr, 'end', '#' * 30
    try:
        session_connections.pop(session_id)
    except KeyError:
        pass

    print '@' * 30

def main():
    svr = gevent.server.StreamServer(('0.0.0.0', 8090), handler)
    svr.serve_forever()

if __name__ == '__main__':
    main()
    
